import os
import requests
import json
from pathlib import Path

# 配置参数
API_URL = "https://api.deepseek.com/v1/chat/completions"
CATEGORIES = ["科学家", "企业家", "文学家", "平凡人"]
OUTPUT_DIR = "biography_materials"

def get_api_key():
    """交互式获取API Key并进行基础验证"""
    while True:
        api_key = input("请输入DeepSeek API Key（输入后按回车）: ").strip()
        if len(api_key) >= 32:  # 简单验证长度
            return api_key
        print("无效的API Key格式！请确认输入的是有效的API Key")

def call_deepseek_api(prompt, api_key):
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "deepseek-chat",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.7,
        "max_tokens": 4000
    }
    try:
        response = requests.post(API_URL, headers=headers, json=data, timeout=30)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"API调用失败: {str(e)}")
        if "401" in str(e):
            print("⚠️ API Key无效或已过期，请重新输入")
        return None

def get_people_list(category, count, api_key):
    prompt = f"""列出{count}位近几年著名的的{category}姓名（确保不重复），中国人为主，外国人这几年特别出名的也可以
    格式为纯文本每行一个姓名，不要编号和其他说明文字"""
    response = call_deepseek_api(prompt, api_key)
    if not response:
        return []
    return [name.strip() for name in response['choices'][0]['message']['content'].split('\n') 
            if name.strip()]

def get_person_details(name, category, api_key):
    prompt = f"""以Markdown格式提供{name}的详细传记资料（{category}类别），包含：
    ## 生平概况
    - 生卒年月
    - 国籍/籍贯
    - 教育背景
    - 主要成就
    
    ## 典型事例（3-5个重要事件）
    - 每个事件包含具体经过和展现的精神品质
    
    {"## 代表作品分析" if category == "文学家" else ""}
    {"- 主要作品及深刻人物分析" if category == "文学家" else ""}
    """
    response = call_deepseek_api(prompt, api_key)
    return response['choices'][0]['message']['content'] if response else None

def save_markdown(category, name, content):
    dir_path = Path(OUTPUT_DIR) / category.replace('/', '_')  # 防止路径问题
    dir_path.mkdir(parents=True, exist_ok=True)
    file_path = dir_path / f"{name.replace('/', '_')}.md"  # 处理特殊字符
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)
        print(f"√ 已保存: {file_path}")
    except IOError as e:
        print(f"文件保存失败: {str(e)}")

def main():
    api_key = get_api_key()
    
    for category in CATEGORIES:
        print(f"\n{'='*30}\n正在处理 [{category}] 类别...\n{'='*30}")
        people = get_people_list(category, 40, api_key)
        
        if not people:
            print(f"⚠️ 未能获取{category}名单，跳过该类别")
            continue
            
        for i, person in enumerate(people[:40], 1):
            print(f"[{i}/{len(people)}] 获取 {person} 的资料...")
            details = get_person_details(person, category, api_key)
            
            if details:
                save_markdown(category, person, details)
            else:
                print(f"⚠️ 跳过 {person}（获取资料失败）")

if __name__ == "__main__":
    print("""DeepSeek人物传记资料采集系统
    (需要有效的DeepSeek API Key)""")
    main()
    print("\n程序执行完毕！输出目录:", os.path.abspath(OUTPUT_DIR))